package com.laotang.第_7_章_堆.力扣_373_查找和最小的;

import java.util.*;

public class _373_find_k_pairs_with_smallest_sums {

    // 大顶堆
    // 时间复杂度：O(mnlogk) --> 超时
    //    1 <= nums1.length, nums2.length <= 10^5
    //    1 <= k <= 10^4
    // 空间复杂度：O(k)
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        // 大顶堆
        PriorityQueue<int[]> pq = new PriorityQueue<>(k, (o1, o2)->{
            return o2[0] + o2[1] - (o1[0] + o1[1]);
        });

        for (int i = 0; i < nums1.length; i++) {
            for (int j = 0; j < nums2.length; j++) {
                if (pq.size() < k) {
                    pq.add(new int[] {nums1[i], nums2[j]});
                } else {
                    int currSum = nums1[i] + nums2[j];
                    int topSum = pq.peek()[0] + pq.peek()[1];
                    if (currSum < topSum) {
                        pq.remove();
                        pq.add(new int[] {nums1[i], nums2[j]});
                    }
                }
            }
        }

        List<List<Integer>> ans = new ArrayList<>();
        while (!pq.isEmpty()) {
            int[] nums = pq.poll();
            ans.add(Arrays.asList(nums[0], nums[1]));
        }

        return ans;
    }


    // 多路归并实现
    // 时间复杂度：O(klog(min(m, k))))
    // 空间复杂度：O(min(m, k))
    public List<List<Integer>> kSmallestPairs2(int[] nums1, int[] nums2, int k) {
        int m = nums1.length, n = nums2.length;
        List<List<Integer>> res = new ArrayList<>();

        // 1. 初始化一个大小为 Math.min(k, m) 的小顶堆
        PriorityQueue<int[]> minHeap = new PriorityQueue<>(Math.min(k, m), new Comparator<int[]>() {
            @Override
            public int compare(int[] o1, int[] o2) {
                // 按照数对和进行排序
                int val1 = nums1[o1[0]] + nums2[o1[1]];
                int val2 = nums1[o2[0]] + nums2[o2[1]];
                return val1 - val2;
            }
        });

        // 将 Math.min(k, m) 个数组的第一个元素值都放入到堆中
        // 这个 for 循环的时间复杂度：O(min(k, m)log(min(k, m)))
        for (int i = 0; i < Math.min(k, m); i++) {
            minHeap.add(new int[]{i, 0});
        }

        //  合并排序
        // 这个 while 循环的时间复杂度：O(klog(min(k, m)))
        while (!minHeap.isEmpty() && k-- > 0) {
            // 拿到堆顶最小的元素
            int[] rec = minHeap.poll();
            int i = rec[0], j = rec[1];

            // 将最小元素放入到结果中
            res.add(Arrays.asList(nums1[i], nums2[j]));

            // 处理这个 元素 所在的数组的下一个元素
            j++;
            // 如果没有下一个元素了，则不做任何处理
            if (j == n) continue;
            // 将下一个元素的位置信息放入到堆中
            minHeap.add(new int[]{i, j});
        }

        return res;
    }
}
